# Dataset: https://www.kaggle.com/gauravduttakiit/covid-19


import pandas
import os
from tools.us_state_abbr import zh_to_en, en_to_zh
from sys import argv
import json


def state_to_init_json(
        geojson_path, population_csv_path, simplified_csv_path, migration_csv_path, recovered_csv_path,
        init_json_path, start_date="2020-01-22"
):
    region_list = []
    geojson_names = [full.split(".")[0] for full in os.listdir(geojson_path)]
    with open(population_csv_path, "r") as population_csv, open(migration_csv_path, "r") as migration_csv, \
            open(simplified_csv_path, "r") as simplified_csv, open(recovered_csv_path, "r") as recovered_csv:
        population_df = pandas.read_csv(population_csv)
        migration_df = pandas.read_csv(migration_csv)
        # 通过实际数据计算感染模型参数
        simplified_df = pandas.read_csv(simplified_csv)
        state_avg_df = simplified_df.groupby("NAME").agg({"Confirmed": "mean", "Deaths": "mean"}).reset_index()
        confirmed_sum_df = simplified_df.groupby("Date").agg({"Confirmed": "sum"}).reset_index()
        recovered_df = pandas.read_csv(recovered_csv)
        for each_state in geojson_names:
            zh_name = each_state
            en_name = zh_to_en[zh_name]
            confirmed_avg_days = int(state_avg_df[state_avg_df.NAME == en_name]["Confirmed"].values[0])
            deaths_avg_days = int(state_avg_df[state_avg_df.NAME == en_name]["Deaths"].values[0])
            confirmed = int(
                simplified_df[(simplified_df.NAME == en_name) & (simplified_df.Date == start_date)]["Confirmed"].values[0]
            )
            dead = int(
                simplified_df[(simplified_df.NAME == en_name) & (simplified_df.Date == start_date)]["Deaths"].values[0]
            )
            confirmed_ratio = \
                confirmed / \
                float(confirmed_sum_df[confirmed_sum_df.Date == start_date]["Confirmed"].values[0])
            recovered = \
                int(recovered_df[recovered_df.Date == start_date]["Recovered"].values[0]) * confirmed_ratio
            infected = confirmed - dead - recovered
            population = int(population_df[population_df.NAME == en_name]["POPESTIMATE2020"].values[0])
            # 数据集中给出的原始数据为一年内的迁入率且为千分制，因此需要进行一定的处理才能转换为所需的迁出率
            export_rate = -migration_df[migration_df.NAME == en_name]["RDOMESTICMIG2019"].values[0] / 1000 / 365
            # 各个地区迁出偏好相等
            zh_names_tuple_except_self = tuple(filter(lambda x: x != zh_name, geojson_names))
            avg_pref_tuple = tuple(1 / len(zh_names_tuple_except_self) for _ in range(len(zh_names_tuple_except_self)))
            dest_pref = dict(zip(zh_names_tuple_except_self, avg_pref_tuple))
            region_list.append({
                "name": zh_name,
                "susceptible": population - infected - dead - recovered,
                "exposed": 0,
                "infected": infected,
                "recovered": recovered,
                "dead": dead,
                "export_rate": export_rate,
                "dest_pref": dest_pref,
                "model_param": {
                    "type": "SEIR",
                    "contact_infected": confirmed_avg_days // 50000,
                    "contact_exposed": confirmed_avg_days // 25000,
                    "alpha": 0.01,
                    "beta": deaths_avg_days / confirmed_avg_days / 80,
                    "lambda_infected": 0.003,
                    "lambda_exposed": 0.0015,
                    "sigma": 1 / 14,
                    "mu": 0.0001
                }
            })
    with open(init_json_path, "w", encoding="utf-8") as init_json:
        json.dump({"regions": region_list}, init_json, indent=4)


def state_to_update_json(geojson_path, population_csv_path, simplified_csv_path, recovered_csv_path, update_json_path):
    update_dict = {}
    geojson_names = [full.split(".")[0] for full in os.listdir(geojson_path)]
    with open(population_csv_path, "r") as population_csv, open(simplified_csv_path, "r") as simplified_csv, \
            open(recovered_csv_path, "r") as recovered_csv:
        population_df = pandas.read_csv(population_csv)
        simplified_df = pandas.read_csv(simplified_csv).sort_values(["Date", "NAME"])
        confirmed_sum_df = simplified_df.groupby("Date").agg({"Confirmed": "sum"}).reset_index()
        recovered_df = pandas.read_csv(recovered_csv)
        date_idx = 0
        cur_date = ""
        cur_date_list = []
        for _, line_dict in simplified_df.iterrows():
            en_name = line_dict["NAME"]
            if not en_to_zh.get(en_name):
                continue
            zh_name = en_to_zh[en_name]
            # 如果时间发生变化则说明当天的数据录入完毕，转入下一天（排序索引为日期与地名）
            if cur_date != line_dict["Date"]:
                update_dict.update({date_idx: cur_date_list})
                cur_date_list = []
                cur_date = line_dict["Date"]
                date_idx += 1
            # Geojson即为仿真程序中的地区名，如果当前地区名不存在对应的geojson文件则跳过
            if zh_name not in geojson_names:
                continue
            # 为了避免出现死亡、感染人口重复计数的问题，需要在population的基础上计算当前易感者
            population = int(population_df[population_df.NAME == en_name]["POPESTIMATE2020"].values[0])
            confirmed_ratio = \
                line_dict["Confirmed"] / \
                float(confirmed_sum_df[confirmed_sum_df.Date == line_dict["Date"]]["Confirmed"].values[0])
            dead = line_dict["Deaths"]
            recovered = \
                int(recovered_df[recovered_df.Date == line_dict["Date"]]["Recovered"].values[0]) * confirmed_ratio
            infected = line_dict["Confirmed"] - line_dict["Deaths"] - recovered
            # 数据集中给出的原始数据为一年内的迁入率且为千分制，因此需要进行一定的处理才能转换为所需的迁出率
            cur_date_list.append({
                "name": zh_name,
                "type": 0,
                "susceptible": population - infected - dead - recovered,
                "infected": infected,
                "dead": dead,
                "recovered": recovered
            })
    with open(update_json_path, "w", encoding="utf-8") as init_json:
        update_dict.pop(0)
        json.dump(update_dict, init_json, indent=4)


def county_csv_to_state_csv(src_csv_path, dst_csv_path, agg_dict):
    with open(src_csv_path, "r") as src_csv:
        src_df = pandas.read_csv(src_csv)
        src_df.rename(columns={"Province/State": "NAME"}).groupby(["NAME", "Date"]).agg(agg_dict).to_csv(dst_csv_path)


def simplify_population_csv(src_csv_path, dst_csv_path):
    with open(src_csv_path, "r") as src_file:
        srf_df = pandas.read_csv(src_file)
        srf_df[["NAME", "POPESTIMATE2020"]].to_csv(dst_csv_path, index=False)


def simplify_migration_csv(src_csv_path, dst_csv_path):
    with open(src_csv_path, "r") as src_file:
        srf_df = pandas.read_csv(src_file)
        srf_df[["NAME", "RDOMESTICMIG2019"]].to_csv(dst_csv_path, index=False)


def simplify_recovered_csv(src_csv_path, dst_csv_path):
    with open(src_csv_path, "r") as src_file:
        srf_df = pandas.read_csv(src_file)
        srf_df[srf_df.Country == "US"][["Date", "Recovered"]].to_csv(dst_csv_path, index=False)


if __name__ == "__main__":
    state_to_init_json(argv[1], argv[2], argv[3], argv[4], argv[5], argv[6], start_date=argv[7])
